/* Copyright 2007-2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Patrick van der Velde
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 * 
 * * Redistributions of source code must retain the above copyright notice, this 
 *   list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright notice, 
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 * * Neither the name of the dnAnalytics Project nor the names of its contributors
 *   may be used to endorse or promote products derived from this software without
 *   specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 * TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

using System;
using System.Collections.Generic;
using dnAnalytics.LinearAlgebra.Solvers.Preconditioners;
using dnAnalytics.Math;

namespace dnAnalytics.LinearAlgebra.Solvers.Iterative
{
    /// <summary>
    /// A Transpose Free Quasi-Minimal Residual (TFQMR) iterative matrix solver.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The TFQMR algorithm was taken from: <br/>
    /// Iterative methods for sparse linear systems.
    /// <br/>
    /// Yousef Saad
    /// <br/>
    /// Algorithm is described in Chapter 7, section 7.4.3, page 219
    /// </para>
    /// <para>
    /// The example code below provides an indication of the possible use of the
    /// solver.
    /// </para>
    /// </remarks>
    /// <include file='../../../../examples/LinearAlgebra/Solvers/Iterative/TFQMR.xml' path='example'/> 
    public sealed class TFQMR : IIterativeSolver
    {
        /// <summary>
        /// The status used if there is no status, i.e. the solver hasn't run yet and there is no
        /// iterator.
        /// </summary>
        private static readonly ICalculationStatus sm_DefaultStatus = new CalculationIndetermined();

        /// <summary>
        /// The preconditioner that will be used. Can be set to <c>null</c>, in which case the default
        /// pre-conditioner will be used.
        /// </summary>
        private IPreConditioner m_Preconditioner;

        /// <summary>
        /// The iterative process controller.
        /// </summary>
        private IIterator m_Iterator;

        /// <summary>
        /// Indicates if the user has stopped the solver.
        /// </summary>
        private bool m_HasBeenStopped;

        /// <summary>
        /// Initializes a new instance of the <c>TFQMR</c> class.
        /// </summary>
        /// <remarks>
        /// When using this constructor the solver will use the <see cref="IIterator"/> with
        /// the standard settings and a default preconditioner.
        /// </remarks>
        public TFQMR() : this(null, null)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>TFQMR</c> class.
        /// </summary>
        /// <remarks>
        /// <para>
        /// When using this constructor the solver will use a default preconditioner.
        /// </para>
        /// <para>
        /// The main advantages of using a user defined <see cref="IIterator"/> are:
        /// <list type="number">
        /// <item>It is possible to set the desired convergence limits.</item>
        /// <item>
        /// It is possible to check the reason for which the solver finished 
        /// the iterative procedure by calling the <see cref="IIterator.Status"/> property.
        /// </item>
        /// </list>
        /// </para>
        /// </remarks>
        /// <param name="iterator">
        ///   The <see cref="IIterator"/> that will be used to monitor the iterative process.
        /// </param>
        public TFQMR(IIterator iterator) : this(null, iterator)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>TFQMR</c> class.
        /// </summary>
        /// <remarks>
        /// When using this constructor the solver will use the <see cref="IIterator"/> with
        /// the standard settings.
        /// </remarks>
        /// <param name="preconditioner">
        /// The <see cref="IPreConditioner"/> that will be used to precondition the matrix equation.
        /// </param>
        public TFQMR(IPreConditioner preconditioner) : this(preconditioner, null)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>TFQMR</c> class.
        /// </summary>
        /// <remarks>
        /// <para>
        /// The main advantages of using a user defined <see cref="IIterator"/> are:
        /// <list type="number">
        /// <item>It is possible to set the desired convergence limits.</item>
        /// <item>
        /// It is possible to check the reason for which the solver finished 
        /// the iterative procedure by calling the <see cref="IIterator.Status"/> property.
        /// </item>
        /// </list>
        /// </para>
        /// </remarks>
        /// <param name="preconditioner">
        /// The <see cref="IPreConditioner"/> that will be used to precondition the matrix equation.
        /// </param>
        /// <param name="iterator">
        /// The <see cref="IIterator"/> that will be used to monitor the iterative process.
        /// </param>
        public TFQMR(IPreConditioner preconditioner, IIterator iterator)
        {
            m_Iterator = iterator;
            m_Preconditioner = preconditioner;
        }

        /// <summary>
        /// Sets the <c>IPreConditioner</c> that will be used to precondition the iterative process.
        /// </summary>
        /// <param name="preconditioner">The preconditioner.</param>
        public void SetPreconditioner(IPreConditioner preconditioner)
        {
            m_Preconditioner = preconditioner;
        }

        /// <summary>
        /// Sets the <c>IIterator</c> that will be used to track the iterative process.
        /// </summary>
        /// <param name="iterator">The iterator.</param>
        public void SetIterator(IIterator iterator)
        {
            m_Iterator = iterator;
        }

        /// <summary>
        /// Gets the status of the iteration once the calculation is finished.
        /// </summary>
        public ICalculationStatus IterationResult
        {
            get 
            { 
                return (m_Iterator != null) ? m_Iterator.Status : sm_DefaultStatus; 
            }
        }

        /// <summary>
        /// Stops the solve process. 
        /// </summary>
        /// <remarks>
        /// Note that it may take an indetermined amount of time for the solver to actually
        /// stop the process.
        /// </remarks>
        public void StopSolve()
        {
            m_HasBeenStopped = true;
        }

        /// <summary>
        /// Solves the matrix equation Ax = b, where A is the coefficient matrix, b is the
        /// solution vector and x is the unknown vector.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="vector">The solution vector, <c>b</c>.</param>
        /// <returns>The result vector, <c>x</c>.</returns>
        public Vector Solve(Matrix matrix, Vector vector)
        {
            if (vector == null)
            {
                throw new ArgumentNullException();
            }

            Vector result = new DenseVector(matrix.Rows);
            Solve(matrix, vector, result);
            return result;
        }

        /// <summary>
        /// Solves the matrix equation Ax = b, where A is the coefficient matrix, b is the
        /// solution vector and x is the unknown vector.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution vector, <c>b</c></param>
        /// <param name="result">The result vector, <c>x</c></param>
        public void Solve(Matrix matrix, Vector input, Vector result)
        {
            // If we were stopped before, we are no longer
            // We're doing this at the start of the method to ensure
            // that we can use these fields immediately.
            {
                m_HasBeenStopped = false;
            }

            // Error checks
            {
                if (matrix == null)
                {
                    throw new ArgumentNullException("matrix");
                }

                if (matrix.Rows != matrix.Columns)
                {
                    throw new MatrixNotSquareException();
                }

                if (input == null)
                {
                    throw new ArgumentNullException("input");
                }

                if (result == null)
                {
                    throw new ArgumentNullException("result");
                }

                if (result.Count != input.Count)
                {
                    throw new NotConformableException();
                }
            }

            // Initialize the solver fields
            {
                // Set the convergence monitor
                if (m_Iterator == null)
                {
                    m_Iterator = Iterator.CreateDefault();
                }

                if (m_Preconditioner == null)
                {
                    m_Preconditioner = new UnitPreconditioner();
                }
                m_Preconditioner.Initialize(matrix);
            }

            var d = new DenseVector(input.Count);
            var r = new DenseVector(input);

            var uOdd = new DenseVector(input.Count);
            var uEven = new DenseVector(input.Count);

            var v = new DenseVector(input.Count);
            var pseudoResiduals = new DenseVector(input);

            var x = new DenseVector(input.Count);
            var yOdd = new DenseVector(input.Count);
            var yEven = new DenseVector(input);

            // Temp vector
            var temp = new DenseVector(input.Count);

            // Initialize
            double startNorm = input.PNorm(2);

            // Define the scalars
            double alpha = 0;
            double eta = 0;
            double theta = 0;

            double tau = startNorm;
            double rho = tau * tau;

            // Calculate the initial values for v
            // M temp = yEven
            m_Preconditioner.Approximate(yEven, temp);
            // v = A temp
            matrix.Multiply(temp, v);

            // Set uOdd
            v.CopyTo(uEven);

            // Start the iteration
            int iterationNumber = 0;
            while (ShouldContinue(iterationNumber, result, input, pseudoResiduals))
            {
                // First part of the step, the even bit
                if (IsEven(iterationNumber))
                {
                    // sigma = (v, r)
                    double sigma = v.DotProduct(r);
                    if (Precision.EqualsWithTolerance(sigma, 0, Constants.UlpsForComparison))
                    {
                        // FAIL HERE
                        //m_Iterator.Fail();
                        m_Iterator.IterationCancelled();
                        break;
                    }
                    // alpha = rho / sigma
                    alpha = rho / sigma;

                    // yOdd = yEven - alpha * v
                    yEven.AddScaledVector(-alpha, v, yOdd);

                    // Solve M temp = yOdd
                    m_Preconditioner.Approximate(yOdd, temp);
                    // uOdd = A temp
                    matrix.Multiply(temp, uOdd);
                }

                // The intermediate step which is equal for both even and
                // odd iteration steps.
                {
                    // Select the correct vector
                    DenseVector uInternal = IsEven(iterationNumber) ? uEven : uOdd;
                    DenseVector yInternal = IsEven(iterationNumber) ? yEven : yOdd;

                    // pseudoResiduals = pseudoResiduals - alpha * uOdd
                    pseudoResiduals.AddScaledVector(-alpha, uInternal);

                    // d = yOdd + theta * theta * eta / alpha * d
                    d.Multiply(theta * theta * eta / alpha, temp);
                    yInternal.Add(temp, d);

                    // theta = ||pseudoResiduals||_2 / tau
                    theta = pseudoResiduals.PNorm(2) / tau;

                    // c = (1 + theta * theta)^-0.5
                    double c = 1 / (System.Math.Sqrt(1 + theta * theta));

                    // tau = tau * theta * c
                    tau *= theta * c;

                    // eta = c^2 * alpha
                    eta = c * c * alpha;

                    // x = x + eta * d
                    x.AddScaledVector(eta, d);

                    // Check convergence and see if we can bail
                    if (!ShouldContinue(iterationNumber, result, input, pseudoResiduals))
                    {
                        // Calculate the real values
                        m_Preconditioner.Approximate(x, result);

                        // Calculate the true residual. Use the temp vector for that
                        // so that we don't pollute the pseudoResidual vector for no
                        // good reason.
                        CalculateTrueResidual(matrix, temp, result, input);

                        // Now recheck the convergence
                        if (!ShouldContinue(iterationNumber, result, input, temp))
                        {
                            // We're all good now.
                            return;
                        }
                    }
                }

                // The odd step
                if (!IsEven(iterationNumber))
                {

                    if (Precision.EqualsWithTolerance(rho, 0, Constants.UlpsForComparison))
                    {
                        // FAIL HERE
                        //m_Iterator.Fail();
                        m_Iterator.IterationCancelled();
                        break;
                    }

                    // rhoNew = pseudoResiduals.DotProduct(r)
                    double rhoNew = pseudoResiduals.DotProduct(r);

                    // beta = rhoNew / rho
                    double beta = rhoNew / rho;

                    // Update rho for the next loop
                    rho = rhoNew;

                    // yOdd = pseudoResiduals + beta * yOdd
                    pseudoResiduals.AddScaledVector(beta, yOdd, yEven);
                    
                    // Solve M temp = yOdd
                    m_Preconditioner.Approximate(yEven, temp);
                    // uOdd = A temp
                    matrix.Multiply(temp, uEven);

                    // v = uEven + beta * (uOdd + beta * v)
                    uOdd.AddScaledVector(beta, v, temp);
                    uEven.AddScaledVector(beta, temp, v);
                }
                
                // Calculate the real values
                m_Preconditioner.Approximate(x, result);

                iterationNumber++;
            }
        }

        // Calculates the true residual of the matrix equation Ax = b according to:
        // residual = b - Ax
        private void CalculateTrueResidual(Matrix matrix, Vector residual, Vector x, Vector b)
        {
            // -Ax = residual
            matrix.Multiply(x, residual);
            residual.Negate();

            // residual + b
            residual.Add(b);
        }

        private bool ShouldContinue(int iterationNumber, Vector result, Vector source, Vector residuals)
        {
            if (m_HasBeenStopped)
            {
                m_Iterator.IterationCancelled();
                return true;
            }

            m_Iterator.DetermineStatus(iterationNumber, result, source, residuals);
            ICalculationStatus status = m_Iterator.Status;

            // We stop if either:
            // - the user has stopped the calculation
            // - the calculation needs to be stopped from a numerical point of view (divergence, convergence etc.)
            return (!status.TerminatesCalculation) && (!m_HasBeenStopped);
        }

        private static bool IsEven(int number)
        {
            return (number % 2 == 0);
        }

        /// <summary>
        /// Solves the matrix equation AX = B, where A is the coefficient matrix, B is the
        /// solution matrix and X is the unknown matrix.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution matrix, <c>B</c>.</param>
        /// <returns>The result matrix, <c>X</c>.</returns>
        public Matrix Solve(Matrix matrix, Matrix input)
        {
            if (matrix == null)
            {
                throw new ArgumentNullException("matrix");
            }
            if (input == null)
            {
                throw new ArgumentNullException("input");
            }
            Matrix result = matrix.CreateMatrix(input.Rows, input.Columns);
            Solve(matrix, input, result);
            return result;
        }

        /// <summary>
        /// Solves the matrix equation AX = B, where A is the coefficient matrix, B is the
        /// solution matrix and X is the unknown matrix.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution matrix, <c>B</c>.</param>
        /// <param name="result">The result matrix, <c>X</c></param>
        public void Solve(Matrix matrix, Matrix input, Matrix result)
        {
            if (matrix == null)
            {
                throw new ArgumentNullException("matrix");
            }
            if (input == null)
            {
                throw new ArgumentNullException("input");
            }
            if (result == null)
            {
                throw new ArgumentNullException("result");
            }

            if (matrix.Rows != input.Rows || input.Rows != result.Rows || input.Columns != result.Columns)
            {
                throw new NotConformableException();
            }

            for (int column = 0; column < input.Columns; column++)
            {
                Vector solution = Solve(matrix, input.GetColumn(column));
                foreach (KeyValuePair<int, double> element in solution.GetIndexedEnumerator())
                {
                    result.ValueAt(element.Key, column, element.Value);
                }
            }
        }
    }
}
